#ifndef _MODEL_H
#define _MODEL_H

#include <vector>
#include <fstream>
#include "biterm.h"
#include "pvec.h"
#include "pmat.h"

using namespace std;

class Model {
public:
  vector<Biterm> bs;

protected:
  int W;				// vocabulary size
  int K0;				// number of common topics
  int K1;				// number of bursty topics
  int n_iter;			// maximum number of iteration of Gibbs Sampling
  int save_step;

  double alpha0;	   // hyperparameters of p(z|t0)	
  double alpha1;	   // hyperparameters of p(z|t1)	
  double beta0;        // hyperparameters of p(w|z,t0)
  double beta1;        // hyperparameters of p(w|z,t1)

  string dfile;			// biterms file
  string dir;			// result dir

  // sample recorders
  Pvec<int> nb_z;	// n(b|z), size K*1
  Pmat<int> nwz;	  // n(w,z), size K*W
  Pmat<double> phi_z;	  // p(w|z), size K*W
  int nt0;				  // number of common biterms
  int nt1;				  // number of bursty biterms

  Pvec<int> nb_c;	// n(b|z), size K*1
  Pmat<int> nwc;	  // n(w,z), size K*W
  Pmat<double> phi_c;	  // p(w|z), size K*W
  
public:
  Model(int K0, int K1, int W, 
		double a0, double a1, double b0, double b1, 
		int n_iter, int save_step,
		string dfile, string dir): 
	K0(K0), K1(K1), W(W), 
	alpha0(a0), alpha1(a1), beta0(b0), beta1(b1),
	n_iter(n_iter), save_step(save_step), 
	dfile(dfile), dir(dir)
	{srand(time(NULL));}
  
  // run estimate and inference procedures
  void run();
  
protected:  
  // intialize memeber varibles and biterms
  void init();		// load from docs
  void load_biterms();

  // update estimate of a biterm
  double update_biterm(Biterm& bi);
  
  // reset topic proportions for biterm b
  void reset_biterm_topic(Biterm& bi);
  
  // assign topic proportions for biterm b
  void assign_biterm_topic(Biterm& bi, int t, int k);
  
  // compute condition distribution p(z|b)
  void compute_pz_b(Biterm& bi, Pvec<double>& p);

  void set_pt_b();

  void save_all();
  void save_pz();
  void save_pw_z();
  void save_pc();
  void save_pw_c();
  void save_btz();
};

#endif
